package unpacket

import (
	"encoding/binary"
	"golang.org/x/net/ipv4"
	"net"
)

func GetPacketHeader(packet []byte) (*ipv4.Header, error) {
	return ipv4.ParseHeader(packet)
}

func ChangeDst(rawPacket []byte, ip string) []byte {
	headerLen := (rawPacket[0] & 0x0f) * 4
	// 修改目标IP地址
	newDstIP := net.ParseIP(ip)
	copy(rawPacket[16:20], newDstIP.To4())

	// 重新计算IP头部校验和
	rawPacket[10] = 0
	rawPacket[11] = 0
	checksum := calculateChecksum(rawPacket[0:headerLen])
	binary.BigEndian.PutUint16(rawPacket[10:12], checksum)
	return rawPacket
}

func ChangeOrigin(rawPacket []byte, ip string) []byte {
	headerLen := (rawPacket[0] & 0x0f) * 4
	// 修改目标IP地址
	newDstIP := net.ParseIP(ip)
	copy(rawPacket[12:16], newDstIP.To4())
	// 重新计算IP头部校验和
	rawPacket[10] = 0
	rawPacket[11] = 0
	checksum := calculateChecksum(rawPacket[0:headerLen])
	binary.BigEndian.PutUint16(rawPacket[10:12], checksum)
	return rawPacket
}

// 计算IP头部校验和
func calculateChecksum(header []byte) uint16 {
	var sum uint32
	for i := 0; i < len(header); i += 2 {
		sum += uint32(header[i])<<8 | uint32(header[i+1])
	}
	sum = (sum >> 16) + (sum & 0xffff)
	sum += sum >> 16
	return uint16(^sum)
}
